import os
import time
import argparse
from glob import glob
import shutil

import mindspore as ms

try:
    from mindformers.tools.logger import logger
except:
     import logging as logger
     logger.getLogger().setLevel(logger.INFO)

from utils import replace_rank_id_in_ckpt_name

ignore_keys_for_resume = ['accu_grads']
ignore_keys_for_eval = ['accu_grads',
                        'scale_sense',
                        'global_step',
                        'adam',
                        'current_iterator_step',
                        'last_overflow_iterator_step',
                        'epoch_num',
                        'step_num',
                        'loss_scale']

def filter_and_save_ckpt(ckpt_path, save_path, mode="resume"):
    if mode == "resume":
        ignore_keys = ignore_keys_for_resume
    elif mode == "eval":
        ignore_keys = ignore_keys_for_eval
    else:
        raise ValueError(f"filter_mode only supports ['resume', 'eval'], but get {mode}")

    checkpoint = ms.load_checkpoint(ckpt_path)
    new_param_list = []
    for name, param in checkpoint.items():
        ignore = False
        for key in ignore_keys:
            if key in name:
                ignore = True
                break
        if not ignore:
            new_param_list.append({"name": name, "data": param})

    ms.save_checkpoint(new_param_list, save_path)
    print(f"process {ckpt_path} finished!")

def main(args):
    checkpoint_dir = args.checkpoint_dir
    save_prefix = args.ckpt_prefix
    process_num = int(args.process_num)
    rank_list = args.rank_list
    filter_mode = args.filter_mode
    logger.info("checkpoint_dir: %s", checkpoint_dir)
    logger.info("save_prefix: %s", save_prefix)
    logger.info("process_num: %d", process_num)
    logger.info("rank_list: %s", rank_list)
    logger.info("filter_mode: %s", filter_mode)

    if checkpoint_dir.endswith("/"):
        checkpoint_dir = checkpoint_dir[:-1]
    new_checkpoint_dir = "checkpoint_filtered"

    rank_dirs = os.listdir(checkpoint_dir)
    new_rank_dirs = []
    for rank_dir in rank_dirs:
        if rank_dir.startswith("rank_"):
            new_rank_dirs.append(rank_dir)
    new_rank_dirs.sort(key=lambda x: int(x.split("rank_")[1]))
    rank_dirs = new_rank_dirs
    logger.info("rank_dirs: %s", rank_dirs)
    total_num = len(rank_dirs)
    logger.info("ckpt total num: %d", total_num)

    if not rank_list:
        rank_id = int(os.getenv("RANK_ID", '0'))
        if total_num % process_num != 0:
            raise ValueError("ckpt_num %d should divided by process_num: %d", total_num, process_num)
        gap_num = total_num // process_num
        start_index = rank_id * gap_num
        end_index = (rank_id + 1) * gap_num
    else:
        rank_dirs = [rank_dirs[idx] for idx in rank_list]
        start_index = 0
        end_index = len(rank_dirs)

    start_time = time.time()
    for rank_dir in rank_dirs[start_index: end_index]:
        logger.info("===========================")
        logger.info(rank_dir)
        dir_rank_id = int(rank_dir.split("rank_")[1])
        prefix = replace_rank_id_in_ckpt_name(save_prefix, dir_rank_id)
        logger.info("prefix: %s", prefix)

        new_rank_dir = os.path.join(new_checkpoint_dir, rank_dir)
        os.makedirs(new_rank_dir, exist_ok=True)

        ckpts = glob(os.path.join(checkpoint_dir, rank_dir, "*.ckpt"))
        new_ckpt = None
        for ckpt in ckpts:
            if os.path.basename(ckpt).startswith(prefix):
                new_ckpt = os.path.join(new_rank_dir, os.path.basename(ckpt))
                filter_and_save_ckpt(ckpt, new_ckpt, filter_mode)
                logger.info("Save %s to %s", ckpt, new_ckpt)
                break
    logger.info("filter finished! cost %f s", time.time() - start_time)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--checkpoint_dir',
        default="./output/checkpoint",
        help='checkpoint dir')
    parser.add_argument(
        '--ckpt_prefix',
        default="llama_7b_rank_0-3_2.ckpt",
        help='ckpt prefix')
    parser.add_argument(
        '--process_num',
        default=1,
        type=int,
        help='process num')
    parser.add_argument(
        '--rank_list',
        default=None,
        type=int,
        nargs='+',
        help='rank list')
    parser.add_argument(
        '--filter_mode',
        default="resume",
        type=str,
        help='resume or eval')
    args = parser.parse_args()

    main(args)
